import pickle  # 替换json模块为pickle
import os
import re
from collections import defaultdict
import jieba
import jieba.analyse


def build_keyword_index(jsonl_path, output_index_path):
    """
    构建术语到知识单元ID的映射索引

    参数:
        jsonl_path: 知识单元JSON Lines文件路径
        output_index_path: 输出的索引文件路径
    """
    # 初始化关键词索引字典，值为知识单元ID的集合
    keyword_index = defaultdict(set)

    # 读取知识单元JSON Lines文件
    with open(jsonl_path, 'r', encoding='utf-8') as f:
        for line in f:
            try:
                # 解析JSON对象
                unit = json.loads(line.strip())
                unit_id = unit['unit_id']
                core_content = unit['core_content']
                source_text = unit['source_text']

                # 合并文本内容以提取更多关键词
                text = core_content + " " + source_text

                # 提取关键词
                # 1. 使用jieba提取关键词（基于TF-IDF）
                tfidf_keywords = jieba.analyse.extract_tags(text, topK=10, withWeight=False)

                # 2. 提取专业术语（大写字母开头的单词和缩写）
                technical_terms = re.findall(r'[A-Z]+[a-zA-Z0-9]*', text)

                # 3. 提取数字相关术语（如参数值、型号等）
                numeric_terms = re.findall(r'[\d\.]+[A-Za-z%°]*', text)

                # 4. 提取中文专业术语（基于词性标注）
                pos_tags = jieba.posseg.cut(text)
                chinese_terms = [word for word, flag in pos_tags if
                                 flag in ['n', 'nr', 'ns', 'nt', 'nz', 'vn'] and len(word) > 1]

                # 合并所有关键词并去重
                all_keywords = list(set(tfidf_keywords + technical_terms + numeric_terms + chinese_terms))

                # 将关键词与知识单元ID建立映射
                for keyword in all_keywords:
                    # 过滤过短的关键词
                    if len(keyword) >= 2 or keyword.isdigit():
                        keyword_index[keyword].add(unit_id)

            except json.JSONDecodeError as e:
                print(f"解析JSON错误: {e}，行内容: {line}")
            except KeyError as e:
                print(f"缺少必要字段: {e}，行内容: {line}")

    # 将索引转换为普通字典（集合转为列表，便于序列化）
    index_dict = {keyword: list(ids) for keyword, ids in keyword_index.items()}

    # 保存索引到pkl文件
    with open(output_index_path, 'wb') as f:  # 注意使用二进制写入模式
        pickle.dump(index_dict, f)

    print(f"关键词索引构建完成！")
    print(f"共提取关键词: {len(index_dict)} 个")
    print(f"索引文件路径: {output_index_path}")

    return index_dict


def load_index(index_path):
    """加载pkl格式的索引文件"""
    with open(index_path, 'rb') as f:  # 注意使用二进制读取模式
        return pickle.load(f)


def search_keyword(index_dict, keyword):
    """
    根据关键词检索知识单元ID

    参数:
        index_dict: 关键词索引字典
        keyword: 要检索的关键词

    返回:
        相关知识单元ID列表，若不存在则返回空列表
    """
    return index_dict.get(keyword, [])


if __name__ == "__main__":
    # 导入json模块（只在需要解析输入JSON时使用）
    import json

    # 获取路径
    script_dir = os.path.dirname(os.path.abspath(__file__))
    project_root = os.path.dirname(script_dir)

    # 知识单元JSON Lines文件路径
    jsonl_file = os.path.join(project_root, "data\knowledge_units.jsonl")

    # 索引文件输出路径（改为pkl后缀）
    index_file = os.path.join(project_root, "data\keyword_index.pkl")

    # 检查知识单元文件是否存在
    if not os.path.exists(jsonl_file):
        raise FileNotFoundError(f"知识单元文件不存在: {jsonl_file}")

    # 构建关键词索引
    keyword_index = build_keyword_index(jsonl_file, index_file)

    # 示例：加载索引并检索
    print("\n示例检索:")
    # 从文件加载索引（实际应用中可以直接使用build_keyword_index返回的结果）
    loaded_index = load_index(index_file)

    sample_keywords = ["TD-LTE", "华为", "功耗", "RSRP"]
    for kw in sample_keywords:
        results = search_keyword(loaded_index, kw)
        print(f"关键词 '{kw}' 找到 {len(results)} 个相关知识单元，前5个ID: {results[:5]}")
